{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Electric Field Conjugation\n",
    "\n",
    "We will implement a basic electric field conjugation with pairwise probing of the electric field. We will use an optical system with both phase and amplitude aberrations, and 2 deformable mirrors for correction.\n",
    "\n",
    "EXPLANATION EFC\n",
    "\n",
    "## With one deformable mirror\n",
    "\n",
    "Let's start by importing the required packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hcipy import *\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll start easy and use an optical system with a single deformable mirror with 32x32 actuators. The dark zone will have an innner working angle of $3 \\lambda/D$ out to an outer working angle of $12\\lambda/D$. As we have only a single deformable mirror, the dark zone will be for $x>1\\lambda/D$ to avoid Hermitian crosstalk between the two sides. The loop gain for the EFC loop is 0.5, so half of the measured electric field will be attempted to be corrected each iteration. We are going to use physical unit in prepartion for using two deformable mirrors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input parameters\n",
    "pupil_diameter = 7e-3 # m\n",
    "wavelength = 700e-9 # m\n",
    "focal_length = 500e-3 # m\n",
    "\n",
    "num_actuators_across = 32\n",
    "actuator_spacing = 1.05 / 32 * pupil_diameter\n",
    "aberration_ptv = 0.02 * wavelength # m\n",
    "\n",
    "epsilon = 1e-9\n",
    "\n",
    "spatial_resolution = focal_length * wavelength / pupil_diameter\n",
    "iwa = 2 * spatial_resolution\n",
    "owa = 12 * spatial_resolution\n",
    "offset = 1 * spatial_resolution\n",
    "\n",
    "efc_loop_gain = 0.5\n",
    "\n",
    "# Create grids\n",
    "pupil_grid = make_pupil_grid(128, pupil_diameter * 1.2)\n",
    "focal_grid = make_focal_grid(2, 16, spatial_resolution=spatial_resolution)\n",
    "prop = FraunhoferPropagator(pupil_grid, focal_grid, focal_length)\n",
    "\n",
    "# Create aperture and dark zone\n",
    "aperture = Field(np.exp(-(pupil_grid.as_('polar').r / (0.5 * pupil_diameter))**30), pupil_grid)\n",
    "\n",
    "dark_zone = circular_aperture(2 * owa)(focal_grid)\n",
    "dark_zone -= circular_aperture(2 * iwa)(focal_grid)\n",
    "dark_zone *= focal_grid.x > offset\n",
    "dark_zone = dark_zone.astype(bool)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create all the optical elements. We'll use a perfect coronagraph as our coronagraph. This can be easily replaced by realistic coronagraphs. Our deformable mirror will have influence functions modelled after Xinetics deformable mirrors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create optical elements\n",
    "coronagraph = PerfectCoronagraph(aperture, order=6)\n",
    "\n",
    "tip_tilt = make_zernike_basis(3, pupil_diameter, pupil_grid, starting_mode=2)\n",
    "aberration = SurfaceAberration(pupil_grid, aberration_ptv, pupil_diameter, remove_modes=tip_tilt, exponent=-3)\n",
    "\n",
    "influence_functions = make_xinetics_influence_functions(pupil_grid, num_actuators_across, actuator_spacing)\n",
    "deformable_mirror = DeformableMirror(influence_functions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Most of the time for EFC we want a function to go from a certain DM pattern to the image in the post-coronagraphic focal plane. Let's create that function now. We'll also include a flag for including aberrations, as we don't want to include the aberrations when calculating the Jacobian matrix of the system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_image(actuators=None, include_aberration=True):\n",
    "    if actuators is not None:\n",
    "        deformable_mirror.actuators = actuators\n",
    "    \n",
    "    wf = Wavefront(aperture, wavelength)\n",
    "    if include_aberration:\n",
    "        wf = aberration(wf)\n",
    "        \n",
    "    img = prop(coronagraph(deformable_mirror(wf)))\n",
    "\n",
    "    return img\n",
    "\n",
    "img_ref = prop(Wavefront(aperture, wavelength)).intensity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's write the function to get the Jacobian matrix. Basically, we poke each actuator and record it's response in post-coronagraphic electric field in the dark zone."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_jacobian_matrix(get_image, dark_zone, num_modes):\n",
    "    responses = []\n",
    "    amps = np.linspace(-epsilon, epsilon, 2)\n",
    "\n",
    "    for i, mode in enumerate(np.eye(num_modes)):\n",
    "        response = 0\n",
    "\n",
    "        for amp in amps:\n",
    "            response += amp * get_image(mode * amp, include_aberration=False).electric_field\n",
    "\n",
    "        response /= np.var(amps)\n",
    "        response = response[dark_zone]\n",
    "\n",
    "        responses.append(np.concatenate((response.real, response.imag)))\n",
    "\n",
    "    jacobian = np.array(responses).T\n",
    "    return jacobian\n",
    "\n",
    "jacobian = get_jacobian_matrix(get_image, dark_zone, len(influence_functions))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can write the EFC routine. We start with nothing on the DM, and record the electric field. For now, we'll use the actual electric field as the estimate; later we'll replace with this a pairwise estimation. We then use the Jacobian from before to retrieve the DM pattern that cancels the estimated electric field. We then put part of the DM correction on the DM and start the next iteration. We'll return the full history of DM patterns, measured electric fields and intensity images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_efc(get_image, dark_zone, num_modes, jacobian, rcond=1e-2):\n",
    "    # Calculate EFC matrix\n",
    "    efc_matrix = inverse_tikhonov(jacobian, rcond)\n",
    "\n",
    "    # Run EFC loop\n",
    "    current_actuators = np.zeros(num_modes)\n",
    "\n",
    "    actuators = []\n",
    "    electric_fields = []\n",
    "    images = []\n",
    "\n",
    "    for i in range(50):\n",
    "        img = get_image(current_actuators)\n",
    "        \n",
    "        electric_field = img.electric_field\n",
    "        image = img.intensity\n",
    "\n",
    "        actuators.append(current_actuators.copy())\n",
    "        electric_fields.append(electric_field)\n",
    "        images.append(image)\n",
    "\n",
    "        x = np.concatenate((electric_field[dark_zone].real, electric_field[dark_zone].imag))\n",
    "        y = efc_matrix.dot(x)\n",
    "\n",
    "        current_actuators -= efc_loop_gain * y\n",
    "        \n",
    "    return actuators, electric_fields, images\n",
    "\n",
    "actuators, electric_fields, images = run_efc(get_image, dark_zone, len(influence_functions), jacobian)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot this as an animation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def make_animation_1dm(actuators, electric_fields, images, dark_zone):\n",
    "    anim = FFMpegWriter('video.mp4', framerate=5)\n",
    "    \n",
    "    num_iterations = len(actuators)\n",
    "    average_contrast = [np.mean(image[dark_zone] / img_ref.max()) for image in images]\n",
    "    \n",
    "    electric_field_norm = mpl.colors.LogNorm(10**-5, 10**(-2.5), True)\n",
    "    \n",
    "    for i in range(num_iterations):\n",
    "        plt.clf()\n",
    "        \n",
    "        plt.subplot(2, 2, 1)\n",
    "        plt.title('Electric field')\n",
    "        electric_field = electric_fields[i] / np.sqrt(img_ref.max())\n",
    "        imshow_field(electric_field, norm=electric_field_norm, grid_units=spatial_resolution)\n",
    "        contour_field(dark_zone, grid_units=spatial_resolution, levels=[0.5], colors='white')\n",
    "    \n",
    "        plt.subplot(2, 2, 2)\n",
    "        plt.title('Intensity image')\n",
    "        imshow_field(np.log10(images[i] / img_ref.max()), grid_units=spatial_resolution, cmap='inferno', vmin=-10, vmax=-5)\n",
    "        plt.colorbar()\n",
    "        contour_field(dark_zone, grid_units=spatial_resolution, levels=[0.5], colors='white')\n",
    "\n",
    "        plt.subplot(2, 2, 3)\n",
    "        deformable_mirror.actuators = actuators[i]\n",
    "        plt.title('DM surface in nm')\n",
    "        imshow_field(deformable_mirror.surface * 1e9, grid_units=pupil_diameter, mask=aperture, cmap='RdBu', vmin=-5, vmax=5)\n",
    "        plt.colorbar()\n",
    "\n",
    "        plt.subplot(2, 2, 4)\n",
    "        plt.title('Average contrast')\n",
    "        plt.plot(range(i), average_contrast[:i], 'o-')\n",
    "        plt.xlim(0, num_iterations)\n",
    "        plt.yscale('log')\n",
    "        plt.ylim(1e-11, 1e-5)\n",
    "        plt.grid(color='0.5')\n",
    "\n",
    "        plt.suptitle('Iteration %d / %d' % (i + 1, num_iterations), fontsize='x-large')\n",
    "    \n",
    "        anim.add_frame()\n",
    "\n",
    "    plt.close()\n",
    "    anim.close()\n",
    "    \n",
    "    return anim\n",
    "\n",
    "make_animation_1dm(actuators, electric_fields, images, dark_zone)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the EFC loop reduces the intensity in the dark zone. The opposite dark zone gets dimmer as well. This is expected as we only introduced phase aberrations. It however doesn't get as dark as the right side, due to frequency folding of speckles outside of the outer working angle.\n",
    "\n",
    "## Add amplitude aberrations\n",
    "\n",
    "Let's now add ampitude aberrations. These are introduced in a telescope by, for example, coating variations and surface errors on optical elementsin a non-conjugate pupil. We will use the latter as an example here and use the built-in `SurfaceAberrationAtDistance` class, which performs a Fresnel propagation to the supplied distance, applies the surface aberration, and performs a Fresnel propagation back to the original plane.\n",
    "\n",
    "Because the `SurfaceAberrationAtDistance` uses a Fresnel propagation, we need to switch to physical units. We redefine our pupil grid and all optical elements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aberration_distance = 100e-3 # m\n",
    "\n",
    "aberration_at_distance = SurfaceAberrationAtDistance(aberration, aberration_distance)\n",
    "\n",
    "aberrated_pupil = aberration_at_distance(Wavefront(aperture, wavelength))\n",
    "\n",
    "plt.subplot(1,2,1)\n",
    "imshow_field(aberrated_pupil.intensity, cmap='inferno')\n",
    "plt.colorbar()\n",
    "plt.subplot(1,2,2)\n",
    "imshow_field(aberrated_pupil.phase, cmap='RdBu', vmin=-0.15, vmax=0.15)\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_image(actuators=None, include_aberration=True):\n",
    "    if actuators is not None:\n",
    "        deformable_mirror.actuators = actuators\n",
    "\n",
    "    wf = Wavefront(aperture, wavelength)\n",
    "    if include_aberration:\n",
    "        wf = aberration_at_distance(wf)\n",
    "        \n",
    "    img = prop(coronagraph(deformable_mirror(wf)))\n",
    "\n",
    "    return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jacobian = get_jacobian_matrix(get_image, dark_zone, len(influence_functions))\n",
    "\n",
    "actuators, electric_fields, images = run_efc(get_image, dark_zone, len(influence_functions), jacobian)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot this as an animation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_animation_1dm(actuators, electric_fields, images, dark_zone)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, the left side only reduces a tiny bit, as all amplitude aberrations are still in the system and we are using only a single deformable mirror for correction.\n",
    "\n",
    "## Adding a second deformable mirror\n",
    "\n",
    "To correct the dark zone on both sides of the star with both phase and amplitude aberrations, we need to use two deformable mirrors, for example, one in the conjugate pupil, and one outside. Let's add the deformable mirrors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "distance_between_dms = 200e-3 # m\n",
    "\n",
    "dm1 = DeformableMirror(influence_functions)\n",
    "dm2 = DeformableMirror(influence_functions)\n",
    "\n",
    "prop_between_dms = FresnelPropagator(pupil_grid, distance_between_dms)\n",
    "\n",
    "dark_zone_twosided = circular_aperture(2 * owa)(focal_grid)\n",
    "dark_zone_twosided -= circular_aperture(2 * iwa)(focal_grid)\n",
    "dark_zone_twosided = dark_zone_twosided.astype('bool')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_image(actuators=None, include_aberration=True):\n",
    "    if actuators is not None:\n",
    "        dm1.actuators = actuators[:len(influence_functions)]\n",
    "        dm2.actuators = actuators[len(influence_functions):]\n",
    "\n",
    "    wf = Wavefront(aperture, wavelength)\n",
    "    if include_aberration:\n",
    "        wf = aberration_at_distance(wf)\n",
    "    \n",
    "    wf_post_dms = prop_between_dms.backward(dm2(prop_between_dms.forward(dm1(wf))))\n",
    "    img = prop(coronagraph(wf_post_dms))\n",
    "\n",
    "    return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jacobian = get_jacobian_matrix(get_image, dark_zone_twosided, 2 * len(influence_functions))\n",
    "\n",
    "actuators, electric_fields, images = run_efc(get_image, dark_zone_twosided, 2 * len(influence_functions), jacobian)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot this as an animation. As we're now using multiple deformable mirrors, we need to create a new plot function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.title('Final focal-plane image')\n",
    "imshow_field(np.log10(images[-1] / img_ref.max()), grid_units=spatial_resolution, cmap='inferno', vmin=-10, vmax=-5)\n",
    "plt.colorbar().set_label('log10 intensity')\n",
    "contour_field(dark_zone_twosided, grid_units=spatial_resolution, levels=[0.5], colors='white')\n",
    "plt.xlabel('x [$\\lambda/D$]')\n",
    "plt.ylabel('y [$\\lambda/D$]')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_animation_2dm(actuators, electric_fields, images, dark_zone):\n",
    "    anim = FFMpegWriter('video.mp4', framerate=5)\n",
    "    \n",
    "    num_iterations = len(actuators)\n",
    "    average_contrast = [np.mean(image[dark_zone] / img_ref.max()) for image in images]\n",
    "    \n",
    "    electric_field_norm = mpl.colors.LogNorm(10**-5, 10**(-2.5), True)\n",
    "    \n",
    "    for i in range(num_iterations):\n",
    "        dm1.actuators = actuators[i][:len(influence_functions)]\n",
    "        dm2.actuators = actuators[i][len(influence_functions):]\n",
    "        \n",
    "        plt.clf()\n",
    "        \n",
    "        plt.subplot(2, 2, 1)\n",
    "        plt.title('Electric field')\n",
    "        electric_field = electric_fields[i] / np.sqrt(img_ref.max())\n",
    "        imshow_field(electric_field, norm=electric_field_norm, grid_units=spatial_resolution)\n",
    "        contour_field(dark_zone, grid_units=spatial_resolution, levels=[0.5], colors='white')\n",
    "    \n",
    "        plt.subplot(2, 2, 2)\n",
    "        plt.title('Intensity image')\n",
    "        imshow_field(np.log10(images[i] / img_ref.max()), grid_units=spatial_resolution, cmap='inferno', vmin=-10, vmax=-5)\n",
    "        plt.colorbar()\n",
    "        contour_field(dark_zone, grid_units=spatial_resolution, levels=[0.5], colors='white')\n",
    "\n",
    "        plt.subplot(2, 3, 4)\n",
    "        plt.title('DM1 surface in nm')\n",
    "        imshow_field(dm1.surface * 1e9, grid_units=pupil_diameter, mask=aperture, cmap='RdBu', vmin=-5, vmax=5)\n",
    "        plt.colorbar()\n",
    "        \n",
    "        plt.subplot(2, 3, 5)\n",
    "        plt.title('DM2 surface in nm')\n",
    "        imshow_field(dm2.surface * 1e9, grid_units=pupil_diameter, mask=aperture, cmap='RdBu', vmin=-5, vmax=5)\n",
    "        plt.colorbar()\n",
    "\n",
    "        plt.subplot(2, 3, 6)\n",
    "        plt.title('Average contrast')\n",
    "        plt.plot(range(i), average_contrast[:i], 'o-')\n",
    "        plt.xlim(0, num_iterations)\n",
    "        plt.yscale('log')\n",
    "        plt.ylim(1e-11, 1e-5)\n",
    "        plt.grid(color='0.5')\n",
    "\n",
    "        plt.suptitle('Iteration %d / %d' % (i + 1, num_iterations), fontsize='x-large')\n",
    "    \n",
    "        anim.add_frame()\n",
    "\n",
    "    plt.close()\n",
    "    anim.close()\n",
    "    \n",
    "    return anim\n",
    "\n",
    "make_animation_2dm(actuators, electric_fields, images, dark_zone_twosided)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.remove('video.mp4')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  },
  "level": "intermediate",
  "thumbnail_figure_index": -1
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
